##      exploits.py
#       
#       Copyright 2009 Hugo Teso <hugo.teso@gmail.com>
#       
#       This program is free software; you can redistribute it and/or modify
#       it under the terms of the GNU General Public License as published by
#       the Free Software Foundation; either version 2 of the License, or
#       (at your option) any later version.
#       
#       This program is distributed in the hope that it will be useful,
#       but WITHOUT ANY WARRANTY; without even the implied warranty of
#       MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#       GNU General Public License for more details.
#       
#       You should have received a copy of the GNU General Public License
#       along with this program; if not, write to the Free Software
#       Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#       MA 02110-1301, USA.

import pygtk
import gtk, gobject

import urllib, tarfile, os, string, csv

class Exploits:

    def __init__(self, config, term_notebook=None):

        # Exploits will be saved on inguma's directory under data/exploits/
        # unless changed on EXPLOITS_DIR variable
        LOCAL_EXPLOIT_DIR = '/data/exploits/'
        self.INGUMA_DIR = os.getcwd()
        self.EXPLOITS_DIR = self.INGUMA_DIR + LOCAL_EXPLOIT_DIR

        self.config = config
        self.term_notebook = term_notebook

        # create a liststore with one string column to use as the model
        self.liststore = gtk.ListStore(int, str, str, str, str, str, str, str)

        self.modelfilter = self.liststore.filter_new()

        # create the TreeView
        self.treeview = gtk.TreeView()

        # create the TreeViewColumns to display the data
        self.treeview.columns = [None]*8
        self.treeview.columns[0] = gtk.TreeViewColumn('No.')
        self.treeview.columns[1] = gtk.TreeViewColumn('Path')
        self.treeview.columns[2] = gtk.TreeViewColumn('Description')
        self.treeview.columns[3] = gtk.TreeViewColumn('Date')
        self.treeview.columns[4] = gtk.TreeViewColumn('Author')
        self.treeview.columns[5] = gtk.TreeViewColumn('Platform')
        self.treeview.columns[6] = gtk.TreeViewColumn('Type')
        self.treeview.columns[7] = gtk.TreeViewColumn('Port')

        # Lets control right click on treeview
        self.treeview.connect('button_press_event', self.on_treeview_button_press_event )

        self.exploits_loaded = 0

        self.show_states = 'All'
        #self.show_states = {'desc':'', 'plat':'', 'type':'', 'port':''}
        self.modelfilter.set_visible_func(self.visible_cb)

        self.treeview.set_model(self.modelfilter)

        # make treeview searchable
        self.treeview.set_enable_search(True)

        # make ui layout
        self.vbox = gtk.VBox()
        self.scrolledwindow = gtk.ScrolledWindow()
        self.bbox = gtk.HButtonBox()
        self.vbox.pack_start(self.bbox, False)
        self.vbox.pack_start(self.scrolledwindow)

        # Textbox for description search
        self.desc = gtk.Entry()
        self.desc.set_text('Text to search')
        self.bbox.pack_start(self.desc)

        # Platform combo filter
        self.platform = gtk.Entry()
        self.platform.set_text('Platform to search')
        self.bbox.pack_start(self.platform)

        # Type combo filter
        self.type = gtk.combo_box_new_text()
        self.type.append_text('Select type:')
        self.type.append_text('remote')
        self.type.append_text('local')
        self.type.append_text('dos')
        self.type.set_active(0)
        self.bbox.pack_start(self.type)

        # Port textbox
        self.textport = gtk.Entry()
        self.textport.set_text('Port to filter')
        self.bbox.pack_start(self.textport)

        # Search Button
        self.search = gtk.Button(stock=gtk.STOCK_FIND)
        self.bbox.pack_start(self.search)

        self.search.connect('clicked', self.check_buttons)

        self.scrolledwindow.add(self.treeview)

    # visibility determined by state matching active toggle buttons
    def visible_cb(self, model, iter):
        #print "Data:", data
        if self.show_states == 'All':
            return True
        else:
            return self.show_states['desc'] in model.get_value(iter, 2) and self.show_states['plat'] in model.get_value(iter, 5) and self.show_states['type'] in model.get_value(iter, 6) and self.show_states['port'] in model.get_value(iter, 7)

    # build list of exploits to show and then refilter
    def check_buttons(self, tb):
        self.show_states = {'desc':'', 'plat':'', 'type':'', 'port':''}

        if self.desc.get_text() != 'Text to search':
            self.show_states['desc'] = self.desc.get_text()
        if self.platform.get_text() != 'Platform to search':
            self.show_states['plat'] = self.platform.get_text()
        if self.get_active_text(self.type) != 'Select type:':
            self.show_states['type'] = self.get_active_text(self.type)
        if self.textport.get_text() != 'Port to filter':
            self.show_states['port'] = self.textport.get_text()

        self.modelfilter.refilter()
        return

    def get_active_text(self, combobox):
        model = combobox.get_model()
        active = combobox.get_active()
        if active < 0:
            return None
        return model[active][0]

    def get_widget(self):
        return self.vbox

    def download_exploits(self, gom):
        """ Downloads and extracts exploits from exploit-db """

        self.gom = gom
        # Move to exploits directory
        try:
        	os.chdir(self.EXPLOITS_DIR)
        except(OSError):
        	os.mkdir(self.EXPLOITS_DIR)
        	os.chdir(self.EXPLOITS_DIR)
        self.gom.echo( "Dir: " + self.EXPLOITS_DIR, False )

        # if exploits dir exists, it's considered an update
        if os.path.exists(self.EXPLOITS_DIR + 'last_date.txt'):
            self.gom.echo( "exploit-db already downloaded, checking for updates", False )
    
        # Fetching exploits from milw0rm
        page = "http://www.exploit-db.com/archive.tar.bz2"
        self.gom.echo( "Downloading " + page, False )
        urllib.urlretrieve(page, "archive.tar.bz2")
    
        # Extract exploits and remove original file
        self.gom.echo( "Extracting files...", False )
        tar = tarfile.open("archive.tar.bz2")
        tar.extractall()
        tar.close()
        os.remove("archive.tar.bz2")
        os.chmod('files.csv', 0644)

        # Let's store actual date
        os.chdir(self.EXPLOITS_DIR)
        self.setDate(self.gom)

        os.chdir(self.EXPLOITS_DIR)
        self.gom.echo( "Operation Complete", False )

        os.chdir(self.INGUMA_DIR)

    def setDate(self, gom):
        """ Stores current date on a file """

        import time
        current_date = time.ctime(time.time())
        f = open('last_date.txt', 'w')
        f.write(current_date)
        f.close()
        self.gom.echo( "Exploits successfully downloaded on" + current_date, False )

    def show_dialog(self):
        """ Pop up dialog to ask for exploits download """

        msg = ("Do you want to download exploits?")
        dlg = gtk.MessageDialog(None, 0, gtk.MESSAGE_QUESTION, gtk.BUTTONS_YES_NO, msg)
        opt = dlg.run()
        dlg.destroy()

        if opt == gtk.RESPONSE_YES:
            return True

        return False

    def check_exploits(self):
        """Check if exploits are already downloaded """

        if os.path.exists(self.EXPLOITS_DIR):
            return True
        else:
            return False

    def load_exploits(self, gom):
        if self.exploits_loaded == 0:
            gom.echo( 'Loading Exploits DDBB...' , False)

            if not os.path.exists(self.EXPLOITS_DIR):
                self.download_exploits(gom)

            # load exploits from csv
            ifile  = open('data/exploits/files.csv', "rb")
            reader = csv.reader(ifile, delimiter=',')
            headerList = reader.next()
    
            # add bug data
            self.states = []
            for line in reader:
                self.liststore.append([ int(line[0]), line[1], line[2], line[3], line[4], line[5], line[6], line[7] ])
    
            self.show_states = 'All'
            self.modelfilter.set_visible_func(self.visible_cb, self.show_states)
    
            self.treeview.set_model(self.modelfilter)
    
            for n in range(8):
                # add columns to treeview
                self.treeview.append_column(self.treeview.columns[n])
                # create a CellRenderers to render the data
                self.treeview.columns[n].cell = gtk.CellRendererText()
                # add the cells to the columns
                self.treeview.columns[n].pack_start(self.treeview.columns[n].cell, True)
                # set the cell attributes to the appropriate liststore column
                self.treeview.columns[n].set_attributes(
                    self.treeview.columns[n].cell, text=n)
    
            # make treeview searchable
            self.treeview.set_search_column(5)
    
            self.exploits_loaded = 1

    def on_treeview_button_press_event(self, treeview, event):
        if event.button == 3:
            x = int(event.x)
            y = int(event.y)
            time = event.time
            pthinfo = treeview.get_path_at_pos(x, y)
            if pthinfo is not None:
                path, col, cellx, celly = pthinfo
                treeview.grab_focus()
                treeview.set_cursor( path, col, 0)

                #print "Right Clicked!"
                treeiter = self.liststore.get_iter(path)
                exploit_path = self.liststore.get_value(treeiter ,1)

                exploit = 'data/exploits/' + exploit_path
                os.chmod(exploit, 0644)
                if self.config.HAS_SOURCEVIEW:
                    import lib.ui.editor as editor
                    editor.main(exploit)
            return True
#        elif event.button == 1 and event.type == gtk.gdk._2BUTTON_PRESS:
#            x = int(event.x)
#            y = int(event.y)
#            pthinfo = treeview.get_path_at_pos(x, y)
#            if pthinfo is not None:
#                path, col, cellx, celly = pthinfo
#                treeview.grab_focus()
#                treeview.set_cursor( path, col, 0)
#
#                #print "Double Left Clicked!"
#                treeiter = self.liststore.get_iter(path)
#                exploit_path = self.liststore.get_value(treeiter ,1)
#
#                exploit = exploit_path.split('/')[-1]
#                exploit_dir = 'data/exploits/' + "/".join(exploit_path.split('/')[0:-1])
#                #print "Exploit: " + exploit
#                self.term_notebook.new_tab(exploit, '. tools/open_exploit ' + exploit_dir)
#
#            return True

